import os

from torchvision import transforms
from torch.utils.data import DataLoader

from dataset import CustomImageDataset
from downdata import load_dataset
from 模型.图像分类模型.其它模型 import *

# 首先下载数据集，这里已经有一个下载数据集的自动代码，你也可以手动下载
load_dataset(root_dir="dataset", dataset_name="红外行人动作分类")


# 设置全局随机种子
torch.manual_seed(123)

device = torch.device("cuda")
# 定义图像转换（可选，根据需要调整）
transform = transforms.Compose([transforms.Resize((224, 224)),  # 调整图像大小以便于显示
    transforms.ToTensor()  # 将图像转换为Tensor
])

# 加载数据集
dataset_train = CustomImageDataset(root_dir='dataset/train', transform=transform)
dataset_test = CustomImageDataset(root_dir='dataset/test', transform=transform)
dataset_val = CustomImageDataset(root_dir='dataset/val', transform=transform)





# 创建一个数据加载器（这里使用shuffle和 ，只是为了测试）
data_loader_train = DataLoader(dataset_train, batch_size=32, shuffle=True)
data_loader_eval = DataLoader(dataset_val, batch_size=1, shuffle=False)
data_loader_test = DataLoader(dataset_test, batch_size=1, shuffle=False)

# 创建模型 创建的模型应该有别人的模型和自己的模型，这样就可以直接比对。
models = []
#图像分类示例
models = [
    'ImageVGG16(num_classes=5, pretrained=True)',
    'ImageResNet50(num_classes=5, pretrained=True)',
    'ImageResNet18(num_classes=5, pretrained=True)',
    'ImageResNet101(num_classes=5, pretrained=True)',
    'ImageDenseNet121(num_classes=5, pretrained=True)',
    'ImageMobileNetV2(num_classes=5, pretrained=True)',
    'ImageResNet50(num_classes=5, pretrained=False)',
]



#视频分类示例：
# models = [
#     'TimeSeriesTransformer(model_name="TimeSeriesTransformer", max_batch=2,feature_dim=768, num_classes=5)',
#     'CNN_GRU(model_name="CNN_GRU",input_size=256, num_classes=5, hidden_size=256, dropout=0.2  , bidirectional=True)',
#     'OurModel(model_name="OurModel", num_classes=5, max_batch=3, pretrained=True)',
#     'XThreeCNN(model_name="XThreeCNN", num_classes=5, max_batch=3, pretrained=True)',
#     'XThreeCNN(model_name="XThreeCNN_NoPretrained", num_classes=5, max_batch=3, pretrained=False)',
#     'SlowFast(model_name="SlowFast_NoPretrained", num_classes= 5, max_batch=3, pretrained=False)',
#     'SlowFast(model_name="SlowFast", num_classes=5, max_batch=3, pretrained=True)',
#     'ThreeResNet(model_name="3DResNet_NoPretrained", num_classes=5, max_batch=2, pretrained=False)',
#     'ThreeResNet(model_name="3DResNet", num_classes=5, max_batch=2, pretrained=True)',
#     'TimesformerModelWrapper(model_name="TimeSformer", max_batch=1, num_classes=5)', ]

"""
在定义模型的时候请注意：
模型输入仅为一个tensor，这个tensor是包含batch_size信息的 图片视频
模型输出的形状仅为一个tensor，这个tensor是包含batch_size信息的分类结果 其shape为(batch_size, num_classes)
"""


for model_str in models:
    model = eval(model_str)
    print("")
    print("开始训练模型：", model_str)
    # 如果 模型参数文件夹下有 model.model_name.pth文件则跳过
    if os.path.exists(f"模型参数/{model_str}.pth"):
        print("模型参数文件已存在，跳过训练")
        continue
    model.my_train(epochs=20, data_loader=data_loader_train, data_loader_eval=data_loader_eval)
    model.my_test(data_loader_eval)
    # 清除显存
    torch.cuda.empty_cache()
